import os

import numpy as np
import cv2
from scipy.constants import point

if __name__ == '__main__':
    # 图片大小最好一样的，不要太大（720P）
    img_path = r'./data/val/images'
    # 关键点检测标签格式，cls,x,y,w,h,x1,y1,v1,x2,y2,v2...或者是cls,x,y,w,h,x1,y1,x2,y2,.....
    # 第一种格式是1+4+n*3的，n个关键点，每个点x坐标，y坐标以及可见性
    # 第二种格式是1+4+n*2的，n个关键点，每个点xy坐标
    # 注意标签都是[0,1]归一化的数据
    label_path = r'./data/val/labels'

    img_names = os.listdir(img_path)
    cls = 0  # 汽车车牌类别

    for file in img_names:
        # 每张图片还是读取到内存中一下，万一有图片大小不一致
        img = cv2.imread(os.path.join(img_path, file))
        H,W = img.shape[:2]
        file_name = os.path.splitext(file)[0]
        info = file_name.split('-')
        # CCPD文件命名格式规范
        box = np.array([item.split('&') for item in info[2].split('_')], dtype=np.float32)
        # 归一化到[0,1]
        box[:,0]/=W
        box[:,1]/=H
        points = np.array([item.split('&') for item in info[3].split('_')], dtype=np.float32)
        # 关键点
        points[:,0]/=W
        points[:,1]/=H
        # 关键点是一行的
        points = points.flatten()

        # 目标框是xcenter,ycenter,w,h格式
        center_x = (box[0,0]+box[1,0])/2
        center_y = (box[0,1]+box[1,1])/2
        w = box[1,0]-box[0,0]
        h = box[1,1]-box[0,1]

        box = np.array([center_x, center_y, w, h])

        result = np.concatenate([box, points])
        # 先写入整数到文件）
        with open(os.path.join(label_path, file_name+'.txt'), 'w') as f:
            f.write(f"{cls} ")

        # 追加写入小数到同一个文件
        with open(os.path.join(label_path, file_name+'.txt'), 'a') as f:
            # 因为delimiter是换行分隔符，所以把数据先reshape一下，每一行数据换行写入文件
            np.savetxt(f, result.reshape(1,-1), fmt='%.8f',delimiter=' ')

    print('ok')